from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate

from ..utils import LOG

API_KEY = "sk-y7DHfp9fzuCxOVm2158638099f9541D3833aB4F4Ed674aCf"
BASE_URL = "https://vip.apiyi.com/v1"  # 此处代理方式，如果是OpenAI官方接口需调整接口地址


class TranslationChain:
    def __init__(self, model_name: str = "", verbose: bool = True):
        template = ("""You are a translation expert, proficient in various languages. \n
            Translates {source_language} to {target_language}.""")

        human_template = "{text}"
        system_message_prompt = SystemMessagePromptTemplate.from_template(template)
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
        chat_prompt_template = ChatPromptTemplate.from_messages(
            [system_message_prompt, human_message_prompt]
        )

        chat = ChatOpenAI(base_url=BASE_URL, api_key=API_KEY, model_name=model_name, temperature=0, verbose=verbose)
        self.chain = LLMChain(llm=chat, prompt=chat_prompt_template, verbose=verbose)

    def run(self, text: str, source_language: str, target_language: str) -> (str, bool):
        result = ""
        try:
            response_dict  = self.chain.invoke({
                "text": text,
                "source_language": source_language,
                "target_language": target_language
            })
            result = response_dict.get("text", "")
        except Exception as e:
            LOG.error(f"An error occurred during translation: {e}")
            return result, False
        return result, True
